{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e697f9e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from py2neo import Graph, Node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2ac145ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "connect successfully\n"
     ]
    }
   ],
   "source": [
    "graph = Graph(\"http://localhost:7474\", auth=(\"neo4j\", \"test\"))\n",
    "if graph:\n",
    "    print(\"connect successfully\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "750d6eb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read data\n",
    "\n",
    "# 7 nodes\n",
    "drugs = []\n",
    "foods = []\n",
    "checks = []\n",
    "departments = []\n",
    "producers = []\n",
    "diseases = []\n",
    "symptoms = []\n",
    "# disease infos\n",
    "disease_infos = []\n",
    "# 11 relationships\n",
    "rels_department = []\n",
    "rels_noteat = []\n",
    "rels_doeat = []\n",
    "rels_recommandeat = []\n",
    "rels_commonddrug = []\n",
    "rels_recommanddrug = []\n",
    "rels_check = []\n",
    "rels_drug_producer = []\n",
    "rels_symptom = []\n",
    "rels_acompany = []\n",
    "rels_category = []\n",
    "\n",
    "data_path = \"data/medical2.json\"\n",
    "for data in open(data_path):\n",
    "    disease_dict = {}\n",
    "    data_json = json.loads(data)\n",
    "    \n",
    "    disease_dict['name'] = disease = data_json['name']\n",
    "    diseases.append(disease)\n",
    "    \n",
    "    disease_dict['desc'] = ''\n",
    "    disease_dict['prevent'] = ''\n",
    "    disease_dict['cause'] = ''\n",
    "    disease_dict['easy_get'] = ''\n",
    "    disease_dict['cure_department'] = ''\n",
    "    disease_dict['cure_way'] = ''\n",
    "    disease_dict['cure_lasttime'] = ''\n",
    "    disease_dict['symptom'] = ''\n",
    "    disease_dict['cured_prob'] = ''\n",
    "    \n",
    "    if 'symptom' in data_json:\n",
    "        # add symptom list\n",
    "        symptoms += data_json['symptom']\n",
    "        # add symptom relationships\n",
    "        for symptom in data_json['symptom']:\n",
    "            rels_symptom.append([disease, symptom])\n",
    "    \n",
    "    if 'acompany' in data_json:\n",
    "        for acompany in data_json['acompany']:\n",
    "            rels_acompany.append([disease, acompany])\n",
    "    \n",
    "    if 'desc' in data_json:\n",
    "        disease_dict['desc'] = data_json['desc']\n",
    "    \n",
    "    if 'prevent' in data_json:\n",
    "        disease_dict['prevent'] = data_json['prevent']\n",
    "    \n",
    "    if 'cause' in data_json:\n",
    "        disease_dict['cause'] = data_json['cause']\n",
    "    \n",
    "    if 'get_prob' in data_json:\n",
    "        disease_dict['get_prob'] = data_json['get_prob']\n",
    "    \n",
    "    if 'easy_get' in data_json:\n",
    "        disease_dict['easy_get'] = data_json['easy_get']\n",
    "    \n",
    "    if 'cure_department' in data_json:\n",
    "        cure_department = data_json['cure_department']\n",
    "        if len(cure_department) == 1:\n",
    "            rels_category.append([disease, cure_department[0]])\n",
    "        elif len(cure_department) == 2:\n",
    "            big = cure_department[0]\n",
    "            small = cure_department[1]\n",
    "            rels_department.append([small, big])\n",
    "            rels_category.append([disease, small])\n",
    "    \n",
    "    if 'cure_way' in data_json:\n",
    "        disease_dict['cure_way'] = data_json['cure_way']\n",
    "        \n",
    "    if 'cure_lasttime' in data_json:\n",
    "        disease_dict['cure_lasttime'] = data_json['cure_lasttime']\n",
    "    \n",
    "    if 'cured_prob' in data_json:\n",
    "        disease_dict['cured_prob'] = data_json['cured_prob']\n",
    "    \n",
    "    if 'common_drug' in data_json:\n",
    "        common_drug = data_json['common_drug']\n",
    "        for drug in common_drug:\n",
    "            rels_commonddrug.append([disease, drug])\n",
    "        drugs += common_drug\n",
    "        \n",
    "    if 'recommand_drug' in data_json:\n",
    "        recommand_drug = data_json['recommand_drug']\n",
    "        drugs += recommand_drug\n",
    "        for drug in recommand_drug:\n",
    "            rels_recommanddrug.append([disease, drug])\n",
    "    \n",
    "    if 'not_eat' in data_json:\n",
    "        not_eat = data_json['not_eat']\n",
    "        for _not in not_eat:\n",
    "            rels_noteat.append([disease, _not])\n",
    "        \n",
    "        foods += not_eat\n",
    "        do_eat = data_json['do_eat']\n",
    "        for _do in do_eat:\n",
    "            rels_doeat.append([disease, _do])\n",
    "        \n",
    "        foods += do_eat\n",
    "        recommand_eat = data_json['recommand_eat']\n",
    "        \n",
    "        for _recommand in recommand_eat:\n",
    "            rels_recommandeat.append([disease, _recommand])\n",
    "        foods += recommand_eat\n",
    "    \n",
    "    if 'check' in data_json:\n",
    "        check = data_json['check']\n",
    "        for _check in check:\n",
    "            rels_check.append([disease, _check])\n",
    "        checks += check\n",
    "    \n",
    "    if 'drug_detail' in data_json:\n",
    "        drug_detail = data_json['drug_detail']\n",
    "        producer = [i.split('(')[0] for i in drug_detail]\n",
    "        rels_drug_producer += [[i.split('(')[0], \n",
    "                                i.split('(')[-1].replace(')', '')] \n",
    "                               for i in drug_detail]\n",
    "        producers += producer\n",
    "    \n",
    "    disease_infos.append(disease_dict)\n",
    "\n",
    "Drugs = set(drugs)\n",
    "Foods = set(foods)\n",
    "Checks = set(checks)\n",
    "Departments = set(departments)\n",
    "Producers = set(producers)\n",
    "Symptoms = set(symptoms)\n",
    "diseases = set(diseases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3d17b71d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create disease nodes\n",
    "for disease_dict in disease_infos:\n",
    "    node = Node(\"Disease\",\n",
    "                name=disease_dict['name'],\n",
    "                desc=disease_dict['desc'],\n",
    "                prevent=disease_dict['prevent'],\n",
    "                cause=disease_dict['cause'],\n",
    "                easy_get=disease_dict['easy_get'],\n",
    "                cure_lasttime=disease_dict['cure_lasttime'],\n",
    "                cure_department=disease_dict['cure_department'],\n",
    "                cure_way=disease_dict['cure_way'],\n",
    "                cured_prob=disease_dict['cured_prob'],\n",
    "               )\n",
    "    graph.create(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e434b020",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create drug nodes\n",
    "for drug in Drugs:\n",
    "    node = Node(\"Drug\", name=drug)\n",
    "    graph.create(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "36750c02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create food nodes\n",
    "for food in Foods:\n",
    "    node = Node(\"Food\", name=food)\n",
    "    graph.create(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "79ffa490",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create check nodes\n",
    "for check in Checks:\n",
    "    node = Node(\"Check\", name=check)\n",
    "    graph.create(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d6b8bf79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create department nodes\n",
    "for department in Departments:\n",
    "    node = Node(\"Department\", name=department)\n",
    "    graph.create(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "01bff317",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create producer nodes\n",
    "for t in Producers:\n",
    "    node = Node(\"Producer\", name=t)\n",
    "    graph.create(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bed05a40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create symptom nodes\n",
    "for t in Symptoms:\n",
    "    node = Node(\"Symptom\", name=t)\n",
    "    graph.create(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0226421d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create relationship\n",
    "def create_relationship(start_node, end_node, edges, rel_type, rel_name):\n",
    "    # deduplicate\n",
    "    set_edges = []\n",
    "    for edge in edges:\n",
    "        set_edges.append('###'.join(edge))\n",
    "    set_edges = set(set_edges)\n",
    "    num_edges = len(set_edges)\n",
    "    for edge in set_edges:\n",
    "        edge = edge.split('###')\n",
    "        p = edge[0]\n",
    "        q = edge[1]\n",
    "        query = \"match(p:%s),(q:%s) where p.name='%s' and q.name='%s' create (p)-[rel:%s{name:'%s'}]->(q)\" % (\n",
    "                start_node, end_node, p, q, rel_type, rel_name)\n",
    "        try:\n",
    "            graph.run(query)\n",
    "        except Exception as e:\n",
    "            print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "388678d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "create_relationship('Disease', 'Food', rels_recommandeat, 'recommand_eat', '推荐食谱')\n",
    "create_relationship('Disease', 'Food', rels_noteat, 'no_eat', '忌吃')\n",
    "create_relationship('Disease', 'Food', rels_doeat, 'do_eat', '宜吃')\n",
    "create_relationship('Department', 'Department', rels_department, 'belongs_to', '属于')\n",
    "create_relationship('Disease', 'Drug', rels_commonddrug, 'common_drug', '常用药品')\n",
    "create_relationship('Producer', 'Drug', rels_drug_producer, 'drugs_of', '生产药品')\n",
    "create_relationship('Disease', 'Drug', rels_recommanddrug, 'recommand_drug', '好评药品')\n",
    "create_relationship('Disease', 'Check', rels_check, 'need_check', '诊断检查')\n",
    "create_relationship('Disease', 'Symptom', rels_symptom, 'has_symptom', '症状')\n",
    "create_relationship('Disease', 'Disease', rels_acompany, 'acompany_with', '并发症')\n",
    "create_relationship('Disease', 'Department', rels_category, 'belongs_to', '所属科室')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
